#!/usr/bin/env python

# for the moment we do without numpy
# import numpy


class SpaceSet(object):
    """A SpaceSet is used to store the relations between spaces
    
        The set contains the mapping rules between the spaces.
    """
    def __init__(self):
        # For the moment we use a stupid implementation
        # All the rules have the form (s1, s2, m)
        # with m a mapping function from s1 to s2
        self.mappings = []
    def add_mapping(self, s1, s2, m):
        self.mappings.append((s1, s2, m))
    def get_mapping(self, s1, s2):
        """Return a mapping from s1 to s2"""
        for r in self.mappings:
            if r[0] is s1 and r[1] is s2:
                return r[2]
        raise Exception("No mapping from %s to %s found" % (s1, s2))
        
    # For the moment we have a single SpaceSet for all the Spaces...
    space_set = None
    
SpaceSet.space_set = SpaceSet()


class Space(object):
    """Base class for any Space"""
    def __init__(self, name = None):
        # this attribute is a map Space -> function
        # that stores the function to transform coordinates from this space to other spaces
        self.set = SpaceSet.space_set
        self.name = name
    def __call__(self, *params):
        """create an Element on the Space"""
        return self.__class__.Element(self, *params)
        
    def get_mapping(self, s):
        return self.set.get_mapping(self, s)
    def add_mapping(self, s, m):
        self.set.add_mapping(self, s, m)
        
    def __repr__(self):
        return self.name
        
    def as_element(self, o):
        """Convert an object to an element of the space.
           
           If the object is already an element of the space do nothing.
           If the object is coordinates then create the new element.
           
           TODO: improve this, it is not that simple
        """
        if isinstance(o, self.__class__.Element):
            return o
        return self(o)
        
    def __pow__(self, n):
        return PowerSpace(self, n)
        
    class Element(object):
        """An element contains its space and its coordinates
        We shouldn't create Element directly but use the 'element' method of the Space instead
        because each Space class defines its own Element Class
        """
        def __init__(self, space, coords):
            self.space = space
            self.__coords = coords
        def coords(self, space = None, mapping = None):
            if space is None or space is self.space:
                return self.__coords
            mapping = mapping or self.space.get_mapping(space)
            return mapping(self).coords()
        def __mul__(self, v):
            # just for the moment
            assert self.space is v.space
            return sum([c1 * c2 for c1,c2 in zip(self.coords(), v.coords())])
        def __repr__(self):
            return repr((self.space, self.__coords))
            

        
class Real(Space):
    """Special Space class for float values"""
    def __init__(self):
        Space.__init__(self, name = 'R')

# The real numbers Space
real = Real()

class Natural(Space):
    def __init__(self):
        Space.__init__(self, name = 'N')

# The natural numbers Space
natural = Natural()

class ProdSpace(Space):
    """Mutli dimensions Space"""
    def __init__(self, spaces, name = None):
        Space.__init__(self, name = name)
        self.spaces = spaces
    def __call__(self, coords):
        return Space.__call__(self, coords)
        
    def __len__(self):
        return len(self.spaces)
    def __getitem__(self, i):
        return self.spaces[i]
        
    def __repr__(self):
        return self.name or ' * '.join([repr(x) for x in self.spaces])
        
    class Element(Space.Element):
        def __init__(self, space, coords):
            Space.Element.__init__(self, space, coords)
            assert len(coords) == len(space)
            

class PowerSpace(ProdSpace):
    """Power space represents Multi-dimensions Space with all dimensions of the same type
    
    That is usefull for very large dimensions space (like discretized mapping)
    """
    def __init__(self, space, n, name = None):
        Space.__init__(self, name = name)
        self.space = space
        self.n = n
    
    def __len__(self):
        return self.n
    def __getitem__(self, i):
        assert i < self.n
        return self.space
    
    def __repr__(self):
        return self.name or '%s**%d' % (repr(self.space), self.n)
    
    class Element(ProdSpace.Element):
        def __init__(self, space, coords):
            ProdSpace.Element.__init__(self, space, coords)
            
class MappingSpace(Space):
    """Special Space that describes Mapping Elements
    @param src   the source Space
    @param dst   the destination Space
    """
    def __init__(self, src, dst, name = None):
        Space.__init__(self, name = name)
        self.src = src
        self.dst = dst
    
    def __repr__(self):
        return "(%s) -> (%s)" % (repr(self.src), repr(self.dst))
        
    def discrete(self, n, vtoi, itod):
        return MappingSpace(self, n, vtoi, itod)
        
    def inverse(self):
        """return the inverse space"""
        return MappingSpace(self.dst, self.src)
        
    class Element(Space.Element):
        def __init__(self, space, func, inverse = None):
            """
            @param func    the function that performs the mapping
            @param inverse give the inverse mapping if the mapping is bijectif
            """
            assert callable(func)
            self.func = func
            self.inverse = inverse and space.inverse().as_element(inverse)
            Space.Element.__init__(self, space, None)
        def __call__(self, element):
            element = self.space.src.as_element(element)
            return self.space.dst(self.func(element.coords(self.space.src)))
        def __repr__(self):
            return repr((self.space, 'function'))
        
        def __mul__(self, e):
            assert isinstance(e, MappingSpace.Element)
            space = MappingSpace(self.space.src, e.space.dst)
            return ComposeElement(space, self, e)
            
class ComposeElement(MappingSpace.Element):
    def __init__(self, space, e1, e2):
        def func(coords):
            return e2(e1(coords).coords()).coords()
        inverse = e1.inverse and e2.inverse and e2.inverse * e1.inverse
        MappingSpace.Element.__init__(self, space, func, inverse)
        
        
class DiscreteMappingSpace(PowerSpace):
    def __init__(self, space, n, vtoi, itod, name = None):
        """
        @param space the Mapping space we want to discretize
        @param n     the number of dimensions of the discrete space
        @param vtoi  a function that maps value to index (space.src -> int)
        @param itod  a function that maps index to discretisation volume (int -> R)
        """
        PowerSpace.__init__(self, space.dst, n, name = name)
        assert callable(vtoi)
        assert callable(itod)
        self.mapping_space = space
        self.vtoi = MappingSpace(space.src, natural).as_element(vtoi)
        self.itod = itod
    
    def __repr__(self):
        return self.name or "D(%s)" % repr(self.mapping_space)
        
    class Element(PowerSpace.Element):
        def __init__(self, space, coords):
            PowerSpace.Element.__init__(self, space, coords)
            self.__deltas = [space.itod(i) for i in range(len(coords))]
        def deltas(self) : return self.__deltas
        
        def __call__(self, element):
            element = self.space.mapping_space.src.as_element(element)
            index = self.space.vtoi(element.coords(self.space.mapping_space.src)).coords()
            return self.space.mapping_space.dst(self.coords()[index])
            

if __name__ == '__main__':

    print '== real space =='
    e = real(1.)
    print 'e =', e  
    
    print '== two real * real spaces, with mapping function =='
    space1 = ProdSpace([real, real], name = 'space1')
    space2 = ProdSpace([real, real], name = 'space2')
    # a mapping function (x,y) -> (x*2, y*2)
    space1.add_mapping(space2, MappingSpace(space1, space2)(lambda x:[x[0]*2, x[1]*2]))
    e = space1([1.,2.])
    print 'e =', e
    print 'cordinates in space1 =', e.coords(space1)
    print 'cordinates in space2 =', e.coords(space2)
    
    print '== power space =='
    space = real**4
    print space
    e = space([1,2,3,4])
    print e
    
    print '== mapping =='
    space = MappingSpace(real, real)
    # a mapping function x -> x**2
    e = space(lambda x:x**2)
    print 'e =', e
    print 'f(2.)=', e(real(2.)).coords()
    print 'f(2.5)=', e(real(2.5)).coords()
    
    print '== discrete mapping (real -> real) with step 1 =='
    # the space values go from 0 to 9 with a step of 1
    # this should be hidden in a special PosMapping class
    space = DiscreteMappingSpace(space, 10, lambda x:int(x), lambda i:1)
    e = space([float(x)**2 for x in range(10)])
    print 'e =', e
    print 'f(2.)=', e(real(2.)).coords()
    print 'f(2.5)=', e(real(2.5)).coords()
    
    print '== more complicated discrete mapping (real * real -> real) =='
    pos_space = ProdSpace([real, real])
    space = MappingSpace(pos_space, real)
    # this should be hidden in a special PosMapping class
    # index from pos function :  i(x,y) = x+2*y
    # pos from imdex function :  x = i%2, y = i/2
    space = DiscreteMappingSpace(space, 4, 
        lambda x:int(x[0])+int(2*x[1]),
        lambda i:1
    )
    # we create a mapping [0,0] -> 0, [0,1] -> 1, [1,0] -> 2, [1,1] -> 3
    e = space(range(4))
    print 'e=', e
    for c in [[0.1,0.2], [1.2,0.4], [1.1,1.4]]:
        print 'f([%g, %g]) = %s' % (c[0], c[1], e(pos_space(c)).coords())
         
    print '== Convolution function conv: F * F -> F, with F: real -> real =='
    # Any function f(x1,x2,...) = y can be represented as an element of a Mapping space
    f_space = MappingSpace(real, real)
    f_space = DiscreteMappingSpace(f_space, 10, lambda x:int(x), lambda i:1)
    conv_space = MappingSpace(ProdSpace([f_space, f_space]), f_space)
    print "conv_space = ", conv_space
    
    def conv_f(e):
        """this function takes as a parameter a list [f1,f2]
           with f1 and f2 elements of 'f_space' (Discrete Mapping R -> R)
           and return the coordinates of a new element of f_space
        """
        f1,f2 = e[0],e[1]
        cf1 = f1.coords()
        cf2 = f2.coords()
        assert len(cf1) == len(cf2)
        n = len(cf1)
        cdelta = f1.deltas()
        # my ad-hoc convolution fuction
        cret = [sum([cf1[i] * cf2[j-i] * cdelta[i] for i in range(j)]) for j in range(n)]
        return cret
            
    conv = conv_space(conv_f)
    
    
    f1 = f_space([x**2 for x in range(10)])
    f2 = f_space([3 * x for x in range(10)])
    
    ret = conv([f1,f2])
    print "f1 =", f1
    print "f2 =", f2
    print "convolution(f1, f2) = ", ret.coords()

    print "== compose mapping =="
    space = MappingSpace(real, real)
    m1 = space(lambda x : 2*x)
    m2 = space(lambda x : x**2)
    m1m2 = m1 * m2
    print m1m2(1)
